package com.eian.boot.common.core.adapter;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.eian.boot.common.id.IdHelper;
import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;

/**
 * @author alex.meng
 * @createTime 2025-09-05 14:50
 * @description TTL增强的MDC上下文管理，支持虚拟线程和普通线程池
 */
public final class TtlMDCAdapter {

    private static final TransmittableThreadLocal<Map<String, String>> CONTEXT = new TransmittableThreadLocal<Map<String, String>>() {
        @Override
        protected Map<String, String> initialValue() {
            return new HashMap<>();
        }

        // 复制父线程的上下文（任务提交时捕获）
        @Override
        public Map<String, String> copy(Map<String, String> parentValue) {
            return parentValue != null ? new HashMap<>(parentValue) : null;
        }

        // 任务执行前：将TTL上下文恢复到MDC（子线程中执行）
        @Override
        protected void beforeExecute() {
            super.beforeExecute();
            Map<String, String> context = get();
            if (context != null) {
                MDC.setContextMap(context);
            }
        }

        // 任务执行后：清理MDC上下文（子线程中执行）
        @Override
        protected void afterExecute() {
            super.afterExecute();
            MDC.clear();
        }
    };

    public static void put(String key, String value) {
        Map<String, String> map = CONTEXT.get();
        if (map == null) {
            map = new HashMap<>();
            CONTEXT.set(map);
        }
        map.put(key, value);
        MDC.put(key, value); // 主线程也同步到MDC
    }

    public static String get(String key) {
        Map<String, String> map = CONTEXT.get();
        return map != null ? map.get(key) : null;
    }

    public static void remove(String key) {
        Map<String, String> map = CONTEXT.get();
        if (map != null) {
            map.remove(key);
        }
        MDC.remove(key);
    }

    public static void clear() {
        CONTEXT.remove();
        MDC.clear();
    }

    // 生成TraceId
    public static String generateTraceId() {
        return IdHelper.simpleUuid();
    }
}